
#include "mainwindow.h"
#include "ui_mainwindow.h"
#include <QMessageBox>
#include <QSqlQuery>
#include <QSqlError>
#include <QCryptographicHash>
#include <QFileDialog>
#include <QFile>
#include <QFileInfo>
#include <QDateTime>
#include <QDebug>
#include "sm3.h"
#include "sm4.h"
#include "sm2.h"

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow)
    , isLoggedIn(false)
{
    ui->setupUi(this);

    connect(ui->loginButton, &QPushButton::clicked, this, &MainWindow::loginUser);
    connect(ui->registerButton, &QPushButton::clicked, this, &MainWindow::registerUser);
    connect(ui->selectFileButton, &QPushButton::clicked, this, &MainWindow::selectFile);
    connect(ui->sendFileButton, &QPushButton::clicked, this, &MainWindow::on_sendFileButton_clicked);
    connect(ui->downloadFileButton, &QPushButton::clicked, this, &MainWindow::downloadFileFromServer);
    connect(ui->auditLogButton, &QPushButton::clicked, this, &MainWindow::viewAuditLog);

    ui->sendFileButton->setEnabled(false);
    ui->downloadFileButton->setEnabled(false);

    createDatabase();
}

MainWindow::~MainWindow()
{
    delete ui;
}

bool MainWindow::createDatabase()
{
    db = QSqlDatabase::addDatabase("QSQLITE");
    db.setDatabaseName("users.db");
    if (!db.open()) return false;

    QSqlQuery query;
    query.exec("CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, password TEXT, salt TEXT, pubkey TEXT)");
    query.exec("CREATE TABLE IF NOT EXISTS auditlog (timestamp TEXT, username TEXT, action TEXT, detail TEXT)");
    return true;
}

QByteArray MainWindow::generateSalt()
{
    return QCryptographicHash::hash(QDateTime::currentDateTime().toString().toUtf8(), QCryptographicHash::Sha256);
}

QByteArray MainWindow::sm3HashWithSalt(const QString &password, const QByteArray &salt)
{
    return sm3(password.toUtf8() + salt);
}

void MainWindow::registerUser()
{
    QString username = ui->usernameLineEdit->text();
    QString password = ui->passwordLineEdit->text();

    if (username.isEmpty() || password.isEmpty()) {
        ui->statusLabel->setText("用户名或密码不能为空");
        return;
    }

    QByteArray salt = generateSalt();
    QByteArray hashedPassword = sm3HashWithSalt(password, salt);

    if (!createSM2Keys(username)) {
        ui->statusLabel->setText("生成SM2密钥失败");
        return;
    }

    QFile pubkeyFile(username + "_pub.pem");
    pubkeyFile.open(QIODevice::ReadOnly);
    QString pubkey = pubkeyFile.readAll();
    pubkeyFile.close();

    QSqlQuery query;
    query.prepare("INSERT INTO users (username, password, salt, pubkey) VALUES (?, ?, ?, ?)");
    query.addBindValue(username);
    query.addBindValue(QString(hashedPassword.toHex()));
    query.addBindValue(QString(salt.toHex()));
    query.addBindValue(pubkey);

    if (!query.exec()) {
        ui->statusLabel->setText("注册失败");
    } else {
        ui->statusLabel->setText("注册成功");
    }
}

void MainWindow::loginUser()
{
    QString username = ui->usernameLineEdit->text();
    QString password = ui->passwordLineEdit->text();

    QSqlQuery query;
    query.prepare("SELECT password, salt FROM users WHERE username = ?");
    query.addBindValue(username);
    if (!query.exec() || !query.next()) {
        ui->statusLabel->setText("用户不存在");
        return;
    }

    QString storedHash = query.value(0).toString();
    QByteArray salt = QByteArray::fromHex(query.value(1).toString().toUtf8());
    QByteArray inputHash = sm3HashWithSalt(password, salt);

    if (!verifySM2Signature(username, password)) {
        ui->statusLabel->setText("SM2签名验证失败");
        return;
    }

    if (storedHash == QString(inputHash.toHex())) {
        ui->statusLabel->setText("登录成功");
        isLoggedIn = true;
        currentUsername = username;
        ui->sendFileButton->setEnabled(true);
        ui->downloadFileButton->setEnabled(true);
        updateLog("登录", "用户登录成功");
    } else {
        ui->statusLabel->setText("密码错误");
    }
}

void MainWindow::selectFile()
{
    selectedFilePath = QFileDialog::getOpenFileName(this, "选择文件");
    ui->statusLabel->setText("已选择文件: " + selectedFilePath);
}

QString MainWindow::sm4Encrypt(const QString &filePath, const QByteArray &key)
{
    QFile file(filePath);
    if (!file.open(QIODevice::ReadOnly)) return "";

    QByteArray data = file.readAll();
    file.close();

    SM4 sm4;
    sm4.setKey(key);
    QByteArray encrypted = sm4.encrypt(data);

    QString newPath = filePath + ".enc";
    QFile encFile(newPath);
    encFile.open(QIODevice::WriteOnly);
    encFile.write(encrypted);
    encFile.close();
    return newPath;
}

QString MainWindow::sm4Decrypt(const QString &filePath, const QByteArray &key)
{
    QFile file(filePath);
    if (!file.open(QIODevice::ReadOnly)) return "";

    QByteArray encrypted = file.readAll();
    file.close();

    SM4 sm4;
    sm4.setKey(key);
    QByteArray decrypted = sm4.decrypt(encrypted);

    QString newPath = filePath + ".dec";
    QFile decFile(newPath);
    decFile.open(QIODevice::WriteOnly);
    decFile.write(decrypted);
    decFile.close();
    return newPath;
}

void MainWindow::on_sendFileButton_clicked()
{
    if (!isLoggedIn || selectedFilePath.isEmpty()) return;

    QByteArray key = "0123456789abcdef"; 
    QString encryptedPath = sm4Encrypt(selectedFilePath, key);
    sendFileToServer(encryptedPath);
    updateLog("上传文件", "上传加密文件: " + encryptedPath);
}

void MainWindow::sendFileToServer(const QString &filePath)
{
    QTcpSocket socket;
    socket.connectToHost("127.0.0.1", 9000);
    if (!socket.waitForConnected(3000)) return;

    QFile file(filePath);
    file.open(QIODevice::ReadOnly);
    socket.write(file.readAll());
    socket.flush();
    socket.waitForBytesWritten();
    socket.disconnectFromHost();
}

void MainWindow::downloadFileFromServer()
{
    if (!isLoggedIn) return;

    QString downloaded = "downloaded.enc";
    QFile file(downloaded);
    file.open(QIODevice::WriteOnly);
    file.write("模拟文件内容");  
    file.close();

    QByteArray key = "0123456789abcdef";
    QString decryptedPath = sm4Decrypt(downloaded, key);
    updateLog("下载文件", "下载并解密文件: " + decryptedPath);
}

bool MainWindow::createSM2Keys(const QString &username)
{
    return sm2_generate_keys(username);
}

bool MainWindow::verifySM2Signature(const QString &username, const QString &password)
{
    QByteArray signature;
    sm2_sign(username + "_pri.pem", password.toUtf8(), signature);
    QFile pubkeyFile(username + "_pub.pem");
    pubkeyFile.open(QIODevice::ReadOnly);
    QByteArray pubkey = pubkeyFile.readAll();
    pubkeyFile.close();
    return sm2_verify(username + "_pub.pem", password.toUtf8(), signature);
}

void MainWindow::updateLog(const QString &action, const QString &detail)
{
    QSqlQuery query;
    query.prepare("INSERT INTO auditlog (timestamp, username, action, detail) VALUES (?, ?, ?, ?)");
    query.addBindValue(QDateTime::currentDateTime().toString(Qt::ISODate));
    query.addBindValue(currentUsername);
    query.addBindValue(action);
    query.addBindValue(detail);
    query.exec();
}

void MainWindow::viewAuditLog()
{
    QSqlQuery query("SELECT * FROM auditlog");
    while (query.next()) {
        QString entry = query.value(0).toString() + " [" + query.value(2).toString() + "] " + query.value(3).toString();
        ui->logTextEdit->append(entry);
    }
}
